package cn.hy.dubbo;

import cn.hy.dubbo.limiter.JudgeResult;
import cn.hy.dubbo.limiter.LimitStrategy;
import cn.hy.dubbo.limiter.TokenBasedLimitStrategy;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author zhaomu
 */
@Activate(
        group = {"consumer", "provider"}
)
public class DubboTpsLimitFilter implements Filter {
    private static final Logger logger = LoggerFactory.getLogger(DubboTpsLimitFilter.class);
    LimitStrategy tpsLimiter = null;

    public DubboTpsLimitFilter() {
    }

    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        if (this.tpsLimiter != null) {
            JudgeResult rs = this.tpsLimiter.isAllowable(invoker.getUrl(), invocation);
            if (!rs.isAllowable()) {
                throw new RpcException(rs.getErrorMsg());
            }
        }
        return invoker.invoke(invocation);
    }

    public void setTokenBasedLimitStrategy(TokenBasedLimitStrategy tokenBasedLimitStrategy) {
        this.tpsLimiter = tokenBasedLimitStrategy;
        logger.info("DubboTpsLimitFilter injection Extension tokenBasedLimitStrategy 完成");
    }
}
